import numpy as np
import os


def gen_golden_data_onehot(num_classes=47, rows=72, cols=123):
    """
    生成二维输入的 one-hot 编码和 golden 数据
    :param num_classes: 类别数量 (depth)
    :param rows: 输入的行数
    :param cols: 输入的列数
    """
    # 生成二维整数索引，范围 [0, num_classes-1]
    input_x = np.random.randint(0, num_classes, size=(rows, cols), dtype=np.int32)

    # 生成 one-hot 编码 [rows, cols, num_classes]
    golden = np.eye(num_classes, dtype=np.int32)[input_x]

    depth_array = np.arange(num_classes, dtype=np.int32)
    # 确保输入输出目录存在
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)

    # 保存
    input_x.tofile("./input/input_x.bin")   # 保存类别索引，shape = [rows, cols]
    golden.tofile("./output/golden.bin")    # 保存 one-hot 向量，shape = [rows, cols, num_classes]


if __name__ == "__main__":
    gen_golden_data_onehot()
